package org.bambrikii.examples.junit.concurrent;

import java.util.concurrent.CountDownLatch;

import org.junit.rules.MethodRule;
import org.junit.runners.model.FrameworkMethod;
import org.junit.runners.model.Statement;

public final class ConcurrentRule implements MethodRule {
	public Statement apply(final Statement statement,
			final FrameworkMethod frameworkMethod, final Object o) {
		return new Statement() {
			@Override
			public void evaluate() throws Throwable {
				final Concurrent concurrent = frameworkMethod
						.getAnnotation(Concurrent.class);
				if (concurrent == null) {
					statement.evaluate();
				} else {
					final String name = frameworkMethod.getName();
					final Thread[] threads = new Thread[concurrent.value()];
					final CountDownLatch go = new CountDownLatch(1);
					final CountDownLatch finished = new CountDownLatch(
							threads.length);
					for (int i = 0; i < threads.length; i++) {
						threads[i] = new Thread(new Runnable() {
							public void run() {
								try {
									go.await();
									statement.evaluate();
								} catch (final InterruptedException e) {
									Thread.currentThread().interrupt();
								} catch (final Throwable throwable) {
									if (throwable instanceof RuntimeException) {
										throw (RuntimeException) throwable;
									}
									if (throwable instanceof Error) {
										throw (Error) throwable;
									}
									final RuntimeException r = new RuntimeException(
											throwable.getMessage(), throwable);
									r.setStackTrace(throwable.getStackTrace());
									throw r;
								} finally {
									finished.countDown();
								}
							}
						}, name + "-Thread-" + i);
						threads[i].start();
					}
					go.countDown();
					finished.await();
				}
			}
		};
	}
}